/*
 * ThreadPool.hpp
 *
 * 一个简单、可复用的 C++17 固定大小线程池实现（头文件）。
 *
 * 说明：
 *  - 声明 ThreadPool 类的接口与必要数据成员。
 *  - 详细注释解释线程池的工作原理及 API。
 *  - 头文件不包含工作线程循环等具体实现，避免重复编译。
 */

#pragma once

#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <future>
#include <functional>
#include <mutex>
#include <queue>
#include <thread>
#include <type_traits>
#include <vector>

class ThreadPool {
public:
    /**
     * @brief 构造一个拥有 threadCount 个工作线程的线程池。
     * @param threadCount 工作线程数量，必须 > 0。
     * @throws std::invalid_argument 当 threadCount == 0 时抛出。
     */
    explicit ThreadPool(std::size_t threadCount);

    /// 禁止拷贝构造与拷贝赋值，防止多个实例共享相同线程资源。
    ThreadPool(const ThreadPool&) = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;

    /// 禁止移动，简化资源管理（如需支持可自行扩展）。
    ThreadPool(ThreadPool&&) = delete;
    ThreadPool& operator=(ThreadPool&&) = delete;

    /**
     * @brief 析构函数，默认调用 shutdown(true) 等待任务完成后关闭。
     */
    ~ThreadPool();

    /**
     * @brief 向线程池提交一个任务，返回 std::future 用于获取结果或异常。
     *
     * @tparam Func 可调用对象类型
     * @tparam Args 参数类型（自动推断）
     * @param func 要执行的可调用对象（支持函数、lambda、函数对象等）
     * @param args 传递给 func 的参数，会被完美转发保持值类别
     *
     * @return std::future<invoke_result_t<Func, Args...>>
     *
     * @note
     *  - 若线程池已停止（显式或析构），调用会抛出 std::runtime_error。
     *  - 线程安全：多个线程同时调用 enqueue 也是安全的。
     */
    template <typename Func, typename... Args>
    auto enqueue(Func&& func, Args&&... args)
        -> std::future<std::invoke_result_t<Func, Args...>>;

    /**
     * @brief 主动关闭线程池。
     *
     * @param wait 是否等待队列中未执行任务完成（默认 true）。
     *             若为 false，将丢弃剩余任务。
     *
     * @note
     *  - shutdown 可安全多次调用，仅第一次生效。
     *  - 调用后将无法再提交新任务。
     */
    void shutdown(bool wait = true);

    /**
     * @brief 获取线程池工作线程数量。
     *
     * @return std::size_t 当前线程数。
     */
    std::size_t size() const noexcept;

private:
    /// @brief 工作线程执行函数，循环获取并执行任务，直至停止。
    void workerLoop(std::size_t workerIndex) noexcept;

private:
    // 工作线程集合。
    std::vector<std::thread> m_workers;

    // 任务队列，存放被包装为 std::function<void()> 的任务。
    std::queue<std::function<void()>> m_tasks;

    // 队列同步原语：互斥量与条件变量。
    std::mutex m_queueMutex;
    std::condition_variable m_condition;

    // 停止标记：true 表示线程池已停止或正在停止。
    std::atomic<bool> m_stop{false};
};

template <typename Func, typename... Args>
auto ThreadPool::enqueue(Func&& func, Args&&... args)
    -> std::future<std::invoke_result_t<Func, Args...>> {
    using ReturnType = std::invoke_result_t<Func, Args...>;

    if (m_stop.load(std::memory_order_acquire)) {
        throw std::runtime_error("Cannot enqueue on stopped ThreadPool.");
    }

    // 使用 shared_ptr 管理 packaged_task 生命周期，确保任务执行前有效。
    auto task = std::make_shared<std::packaged_task<ReturnType()>>(
        [func = std::forward<Func>(func), ... params = std::forward<Args>(args)]() mutable {
            return std::invoke(std::move(func), std::move(params)...);
        }
    );

    std::future<ReturnType> resultFuture = task->get_future();

    {
        std::lock_guard<std::mutex> lock(m_queueMutex);
        if (m_stop.load(std::memory_order_relaxed)) {
            throw std::runtime_error("Cannot enqueue on stopped ThreadPool.");
        }
        m_tasks.emplace([task]() { (*task)(); });
    }

    m_condition.notify_one();
    return resultFuture;
}